package com.yuan.aiagent.rag;

import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.reader.markdown.MarkdownDocumentReader;
import org.springframework.ai.reader.markdown.config.MarkdownDocumentReaderConfig;
import org.springframework.core.io.Resource;
import org.springframework.core.io.support.ResourcePatternResolver;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.ArrayList;
import java.util.Collection;
import java.util.List;

@Component
@Slf4j
public class AppDocumentLoader<Document> {
    private ResourcePatternResolver resolvers;

    public AppDocumentLoader(ResourcePatternResolver resolver) {
        this.resolvers = resolver;
    }

    public List<Document> loadMarkdowns() {
        List<Document> allDocuments = new ArrayList<>();
        try {
            // 加载document目录下的恋爱知识库文档
            Resource[] resources = resolvers.getResources("classpath:document/*.md");
            for (Resource resource : resources) {
                String fileName = resource.getFilename();
                String statue = "恋爱";
                MarkdownDocumentReaderConfig config = MarkdownDocumentReaderConfig.builder()
                        .withHorizontalRuleCreateDocument(true)
                        .withIncludeCodeBlock(false)
                        .withIncludeBlockquote(false)
                        .withAdditionalMetadata("filename", fileName)
                        .withAdditionalMetadata("statue", statue)
                        .withAdditionalMetadata("type", "love")
                        .build();
                MarkdownDocumentReader reader = new MarkdownDocumentReader(resource, config);
                allDocuments.addAll((Collection<? extends Document>) reader.get());
            }
            
            // 加载AFSIM目录下的AFSIM知识库文档
            loadAfsimDocuments(allDocuments);
            
        } catch (IOException e) {
            //log.error("Markdown 文档加载失败", e);
        }
        return allDocuments;
    }
    
    private void loadAfsimDocuments(List<Document> allDocuments) {
        try {
            Resource[] afsimResources = resolvers.getResources("classpath:AFSIM/*.md");
            for (Resource resource : afsimResources) {
                String fileName = resource.getFilename();
                String docType = getAfsimDocumentType(fileName);
                
                MarkdownDocumentReaderConfig config = MarkdownDocumentReaderConfig.builder()
                        .withHorizontalRuleCreateDocument(true)
                        .withIncludeCodeBlock(true) // 包含代码块，这对AFSIM模板很重要
                        .withIncludeBlockquote(true)
                        .withAdditionalMetadata("filename", fileName)
                        .withAdditionalMetadata("statue", "AFSIM")
                        .withAdditionalMetadata("type", docType)
                        .withAdditionalMetadata("domain", "AFSIM")
                        .build();
                MarkdownDocumentReader reader = new MarkdownDocumentReader(resource, config);
                allDocuments.addAll((Collection<? extends Document>) reader.get());
            }
        } catch (IOException e) {
            //log.error("AFSIM Markdown 文档加载失败", e);
        }
    }
    
    private String getAfsimDocumentType(String fileName) {
        if (fileName.contains("基础概念")) {
            return "concept";
        } else if (fileName.contains("代码示例")) {
            return "code_example";
        } else if (fileName.contains("应用场景")) {
            return "best_practice";
        } else if (fileName.contains("模型参数")) {
            return "model_params";
        } else {
            return "general";
        }
    }
}